package com.simen.hellobytes.encrypt;

import com.simen.hellobytes.exception.EncryptException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.*;

/**
 * @author Created by Simen.
 * @date 创建日期 2018/9/18 16:59
 * @modify 修改者 Simen
 */
public class RSAUtil {

    public static final String ALGORITHM_RSA = "RSA";

    /**
     * @param keySize
     * @throws NoSuchAlgorithmException
     */
    public static RSAKey[] generateRsaKeyPair(int keySize) throws NoSuchAlgorithmException, InvalidKeySpecException, IOException {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance(ALGORITHM_RSA);

        kpg.initialize(keySize);

        KeyPair keyPair = kpg.genKeyPair();

        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();

        return new RSAKey[]{publicKey, privateKey};
    }

    public static PublicKey getPublicKeyFromCer(InputStream cerFileInStream) {
        CertificateFactory certificatefactory = null;
        try {
            certificatefactory = CertificateFactory.getInstance("X.509");

            X509Certificate Cert = (X509Certificate) certificatefactory.generateCertificate(cerFileInStream);

            return Cert.getPublicKey();
        } catch (CertificateException e) {
            e.printStackTrace();
        }

        return null;
    }

    public static PublicKey getPublicKey(byte[] encodedKey) throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);

        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(encodedKey);

        return keyFactory.generatePublic(x509KeySpec);
    }

    public static PrivateKey getPrivateKey(byte[] encodedKey) throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);

        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encodedKey);

        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * RSA加解密:
     * 1024位的证书，加密时最大支持117个字节，解密时为128；
     * 2048位的证书，加密时最大支持245个字节，解密时为256;
     * (加密时支持的最大字节数：证书位数/8 -11)
     *
     * @param srcData   要加密的数据
     * @param publicKey 公钥
     * @return 返回密文信息
     */
    public static byte[] encryptByPublicKey(byte[] srcData, Key publicKey) throws EncryptException {
        ByteArrayOutputStream pubOutStream = null;

        try {
            Cipher cipher = Cipher.getInstance(KeyFactory.getInstance(ALGORITHM_RSA).getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);

            int baseOffset = getPublicKeyLength(publicKey) / 8 - 11;

            pubOutStream = new ByteArrayOutputStream();

            separateCipher(srcData, baseOffset, cipher, pubOutStream);

            return pubOutStream.toByteArray();
        } catch (InvalidKeyException e) {
            throw new EncryptException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        } catch (NoSuchPaddingException e) {
            throw new EncryptException(e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptException(e);
        } catch (BadPaddingException e) {
            throw new EncryptException(e);
        } finally {
            try {
                if (pubOutStream != null) {
                    pubOutStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static byte[] decryptByPrivateKey(byte[] srcData, Key privateKey) throws EncryptException {
        ByteArrayOutputStream priOutStream = null;
        try {
            Cipher cipher = Cipher.getInstance(KeyFactory.getInstance(ALGORITHM_RSA).getAlgorithm());

            cipher.init(Cipher.DECRYPT_MODE, privateKey);

            priOutStream = new ByteArrayOutputStream();

            separateCipher(srcData, getPrivateKeyLength(privateKey) / 8, cipher, priOutStream);

            return priOutStream.toByteArray();
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        } catch (NoSuchPaddingException e) {
            throw new EncryptException(e);
        } catch (InvalidKeyException e) {
            throw new EncryptException(e);
        } catch (BadPaddingException e) {
            throw new EncryptException(e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptException(e);
        } finally {
            try {
                if (priOutStream != null) {
                    priOutStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static byte[] encryptByPrivateKey(byte[] srcData, PrivateKey privateKey) throws EncryptException {
        ByteArrayOutputStream priOutStream = null;

        try {
            Cipher cipher = Cipher.getInstance(KeyFactory.getInstance(ALGORITHM_RSA).getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, privateKey);

            int baseOffset = getPrivateKeyLength(privateKey) / 8 - 11;

            priOutStream = new ByteArrayOutputStream();

            separateCipher(srcData, baseOffset, cipher, priOutStream);

            return priOutStream.toByteArray();
        } catch (InvalidKeyException e) {
            throw new EncryptException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        } catch (NoSuchPaddingException e) {
            throw new EncryptException(e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptException(e);
        } catch (BadPaddingException e) {
            throw new EncryptException(e);
        } finally {
            try {
                if (priOutStream != null) {
                    priOutStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static byte[] decryptByPublicKey(byte[] srcData, Key publicKey) throws EncryptException {
        ByteArrayOutputStream pubOutStream = null;
        try {
            Cipher cipher = Cipher.getInstance(KeyFactory.getInstance(ALGORITHM_RSA).getAlgorithm());

            cipher.init(Cipher.DECRYPT_MODE, publicKey);

            pubOutStream = new ByteArrayOutputStream();

            separateCipher(srcData, getPublicKeyLength(publicKey) / 8, cipher, pubOutStream);

            return pubOutStream.toByteArray();
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        } catch (NoSuchPaddingException e) {
            throw new EncryptException(e);
        } catch (InvalidKeyException e) {
            throw new EncryptException(e);
        } catch (BadPaddingException e) {
            throw new EncryptException(e);
        } catch (IllegalBlockSizeException e) {
            throw new EncryptException(e);
        } finally {
            try {
                if (pubOutStream != null) {
                    pubOutStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 分段加解密实现
     */
    private static void separateCipher(byte[] srcData, int baseOffset, Cipher cipher, ByteArrayOutputStream outputStream)
            throws BadPaddingException, IllegalBlockSizeException {
        int inputLen = srcData.length;
        int offSet = 0;

        for (int i = 0; inputLen - offSet > 0; offSet = i * baseOffset) {
            byte[] cache;

            if (inputLen - offSet > baseOffset) {
                cache = cipher.doFinal(srcData, offSet, baseOffset);
            } else {
                cache = cipher.doFinal(srcData, offSet, inputLen - offSet);
            }

            outputStream.write(cache, 0, cache.length);

            ++i;
        }
    }

    /**
     * 获取公钥长度(1024,2048....)
     *
     * @param publicKey
     * @return
     */
    public static int getPublicKeyLength(Key publicKey) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);

            RSAPublicKeySpec keySpec = keyFactory.getKeySpec(publicKey, RSAPublicKeySpec.class);

            BigInteger prime = keySpec.getModulus();

            return prime.toString(2).length();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        }

        return 0;
    }

    /**
     * 获取私钥长度(1024,2048....)
     *
     * @param privateKey
     * @return
     */
    public static int getPrivateKeyLength(Key privateKey) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);

            RSAPrivateKeySpec keySpec = keyFactory.getKeySpec(privateKey, RSAPrivateKeySpec.class);

            BigInteger prime = keySpec.getModulus();

            return prime.toString(2).length();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        }

        return 0;
    }

}
